# This code is referenced from 
# https://github.com/facebookresearch/astmt/
# 
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
# 
# License: Attribution-NonCommercial 4.0 International

import warnings
import cv2
import os.path
import numpy as np
import glob
import torch
import json
import scipy.io as sio

from termcolor import colored


class HeightMeter(object):
    def __init__(self):
        self.total_rmses = []
        self.total_zncc = []

    @torch.no_grad()
    def update(self, pred, gt):
        pred, gt = pred.squeeze(), gt.squeeze()

        # Determine valid mask
        for i in range(gt.shape[0]):
            mask = gt[i] >= 0
            n_valid = mask.sum()  # Valid pixels per image

            # Only positive depth values are possible
            # pred[pred < 1e-9] = 1e-9
            #
            # # Per pixel rmse and log-rmse.
            # log_rmse_tmp = np.power(np.log(gt) - np.log(pred), 2)
            # log_rmse_tmp = log_rmse_tmp[mask]
            # self.total_log_rmses += log_rmse_tmp.sum()

            rmse_tmp = np.power(gt[i] - pred[i], 2)
            rmse_tmp = rmse_tmp[mask]
            self.total_rmses.append(np.sqrt(rmse_tmp.sum()/n_valid))

            w1 = gt[i] - gt[i].mean()
            w2 = pred[i] - pred[i].mean()
            denom = np.sqrt(np.sum(w1 ** 2) * np.sum(w2 ** 2))
            if denom < 1e-10:
                self.total_zncc.append(0)
            else:
                self.total_zncc.append(np.sum(w1 * w2) / denom)

    def reset(self):
        self.total_rmses = []
        self.total_zncc = []

    def get_score(self, is_val=False, verbose=True):
        eval_result = dict()
        eval_result['rmse'] = np.asarray(self.total_rmses).mean()
        eval_result['zncc'] = np.asarray(self.total_zncc).mean()
        # if verbose:
        if is_val:
            print(colored('Results for height prediction', 'yellow'))
            for x in eval_result:
                # spaces = ''
                # for j in range(0, 15 - len(x)):
                #     spaces += ' '
                print(colored('{0:s}:{1:.4f}'.format(x, eval_result[x]), 'yellow'), end='   ')
            print()
            print()
        else:
            print('Results for height prediction')
            for x in eval_result:
                # spaces = ''
                # for j in range(0, 15 - len(x)):
                #     spaces += ' '
                print('{0:s}:{1:.4f}'.format(x, eval_result[x]), end='   ')
            print()
            print()

        return eval_result


Vaihingen_CATEGORY_NAMES = ['ImSurf', 'Building', 'LowVeg',
                            'Tree', 'Car', 'Clutter']


class SemsegMeter(object):
    def __init__(self, database):
        self.database = database
        if database == 'Vaihingen':
            self.num_class = 6
            self.cat_names = Vaihingen_CATEGORY_NAMES
        else:
            raise NotImplementedError

        self.confusion_matrix = np.zeros((self.num_class,) * 2)

    @torch.no_grad()
    def update(self, pred, gt):
        pred = pred.squeeze()
        gt = gt.squeeze()
        self.confusion_matrix += self._generate_matrix(gt, pred)

    def reset(self):
        self.confusion_matrix = np.zeros((self.num_class,) * 2)

    def get_score(self, is_val=False, verbose=True):
        eval_result = dict()
        if self.database == 'vaihingen':
            eval_result['mIoU'] = np.mean(self.IoU()[:-1])
            eval_result['F1'] = np.mean(self.F1()[:-1])
        else:
            eval_result['mIoU'] = np.mean(self.IoU())
            eval_result['F1'] = np.mean(self.F1())
        eval_result['OA'] = np.mean(self.OA())

        text_color = 'yellow' if is_val else None
        print(colored('Semantic Segmentation mIoU: {0:.4f}\n'.format(eval_result['mIoU'], text_color)))
        print(colored('Semantic Segmentation F1: {0:.4f}\n'.format(np.mean(eval_result['F1']), text_color)))
        print(colored('Semantic Segmentation OA: {0:.4f}\n'.format(np.mean(eval_result['OA']), text_color)))

        if verbose:
            iou_value = {}
            iou_per_class = self.IoU()
            for class_name, iou in zip(self.cat_names, iou_per_class):
                iou_value[class_name] = iou
                spaces = ''
                for j in range(0, 20 - len(class_name)):
                    spaces += ' '
                print(colored('{0:s}{1:s}{2:.4f}'.format(class_name, spaces, iou), text_color))

        return eval_result

    def get_tp_fp_tn_fn(self):
        tp = np.diag(self.confusion_matrix)
        fp = self.confusion_matrix.sum(axis=0) - np.diag(self.confusion_matrix)
        fn = self.confusion_matrix.sum(axis=1) - np.diag(self.confusion_matrix)
        tn = np.diag(self.confusion_matrix).sum() - np.diag(self.confusion_matrix)
        return tp, fp, tn, fn

    def Precision(self):
        tp, fp, tn, fn = self.get_tp_fp_tn_fn()
        precision = tp / (tp + fp)
        return precision

    def Recall(self):
        tp, fp, tn, fn = self.get_tp_fp_tn_fn()
        recall = tp / (tp + fn)
        return recall

    def F1(self):
        tp, fp, tn, fn = self.get_tp_fp_tn_fn()
        Precision = tp / (tp + fp)
        Recall = tp / (tp + fn)
        F1 = (2.0 * Precision * Recall) / (Precision + Recall)
        return F1

    def OA(self):
        OA = np.diag(self.confusion_matrix).sum() / (self.confusion_matrix.sum() + self.eps)
        return OA

    def IoU(self):
        tp, fp, tn, fn = self.get_tp_fp_tn_fn()
        IoU = tp / (tp + fn + fp)
        return IoU

    def Dice(self):
        tp, fp, tn, fn = self.get_tp_fp_tn_fn()
        Dice = 2 * tp / ((tp + fp) + (tp + fn))
        return Dice

    def Pixel_Accuracy_Class(self):
        #         TP                                  TP+FP
        Acc = np.diag(self.confusion_matrix) / (self.confusion_matrix.sum(axis=0) + self.eps)
        return Acc

    def Frequency_Weighted_Intersection_over_Union(self):
        freq = np.sum(self.confusion_matrix, axis=1) / (np.sum(self.confusion_matrix) + self.eps)
        iou = self.IoU()
        FWIoU = (freq[freq > 0] * iou[freq > 0]).sum()
        return FWIoU

    def _generate_matrix(self, gt_image, pre_image):
        mask = (gt_image >= 0) & (gt_image < self.num_class)
        label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]
        count = np.bincount(label, minlength=self.num_class ** 2)
        confusion_matrix = count.reshape(self.num_class, self.num_class)
        return confusion_matrix

